import random
import os
import argparse
from glob import glob
import shutil
class_name={"normal": 0, "anthrax": 1, "leafhopper": 2}
train_val_percent = 0.9
train_percent = 0.8
file_train_value = None
file_train = None
file_val = None
file_test = None

def get_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--anthrax_text_path', default=r'D:\project\code\tea\pest-and-disease-identification\picture\label-2022-05-05-anthrax/',
                        type=str, help='input xml file ')
    parser.add_argument('--leafhopper_txt_path', default=r"D:\project\code\tea\pest-and-disease-identification\picture\label-2022-04-02-leafhopper/",
                        type=str, help='output txt file')
    parser.add_argument('--normal_txt_path',
                        default=r"D:\project\code\tea\pest-and-disease-identification\picture\label-2022-05-15-normal/",
                        type=str, help='output txt file')
    parser.add_argument('--anthrax_img_path',
                        default=r"D:\project\code\tea\pest-and-disease-identification\picture\2022-05-05-Anthrax/",
                        type=str, help='output txt file')
    parser.add_argument('--leafhopper_img_path',
                        default=r"D:\project\code\tea\pest-and-disease-identification\picture\2022-04-02-Small green leafhopper/",
                        type=str, help='output txt file')
    parser.add_argument('--normal_img_path',
                        default=r"D:\project\code\tea\pest-and-disease-identification\picture\2022-05-15-normal/",
                        type=str, help='output txt file')
    opt = parser.parse_args()
    return opt

def generate_random_file(filepath, classname):

    src_file_list = glob(filepath + classname+'*.txt')  # glob获得路径下所有文件，可根据需要修改
    total_txt = len(src_file_list)
    list_index = range(total_txt)
    train_value_num = int(total_txt * train_val_percent)
    train_num = int(train_value_num*train_percent)
    train_val_index = random.sample(list_index, train_value_num)
    train_index = random.sample(list_index, train_num)

    return train_val_index, train_index, list_index, src_file_list

def create_txt_file(path):
    global file_train_value
    global file_train
    global file_test
    global file_val
    file_train_value = open(path +'/train_val.txt', 'a+')
    file_train = open(path + '/train.txt', 'a+')
    file_test = open(path + '/test.txt', 'a+')
    file_val = open(path + '/val.txt', 'a+')

    # file_train = open(path + '/train.txt', 'r')
    # txtlist = file_train.readlines()
    # txtlist = file_train.readlines()

def split_data(filepath, classname):
    train_val_index, train_index, list_index, src_file= generate_random_file(filepath, classname)
    for i in list_index:
        data_name = src_file[i][:-4]+'\n'
        if i in train_val_index:
            file_train_value.write(data_name)
            if i in train_index:
                file_train.write(data_name)
            else:
                file_val.write(data_name)
        else:
            file_test.write(data_name)

def copy_flie(path, opt):
    file_train_value = open(path + '/train_val.txt', 'r')
    file_train = open(path + '/train.txt', 'r')
    file_test = open(path + '/test.txt', 'r')
    file_val = open(path + '/val.txt', 'r')
    anthrax_img_path = opt.anthrax_img_path
    leafhopper_img_path = opt.leafhopper_img_path
    normal_img_path = opt.normal_img_path

    txtlist = file_train.readlines()
    for filepath in txtlist:
        oneline = filepath.strip()
        filename = oneline.split('\\')[-1]
        shutil.copy(oneline+".txt", os.getcwd()+'/labels/train/'+str(filename)+".txt")
        if(filename.split('_')[0] == "normal"):
            shutil.copy(normal_img_path + filename +  ".jpg", os.getcwd() + '/images/train/' + str(filename) + ".jpg")
        if (filename.split('_')[0] == "anthrax"):
            shutil.copy(anthrax_img_path + filename + ".jpg", os.getcwd() + '/images/train/' + str(filename) + ".jpg")
        if (filename.split('_')[0] == "leafhopper"):
            shutil.copy(leafhopper_img_path + filename + ".jpg", os.getcwd() + '/images/train/' + str(filename) + ".jpg")

    txtlist = file_val.readlines()
    for filepath in txtlist:
        oneline = filepath.strip()
        filename = oneline.split('\\')[-1]
        shutil.copy(oneline + ".txt", os.getcwd() + '/labels/value/' + str(filename) + ".txt")

        if (filename.split('_')[0] == "normal"):
            shutil.copy(normal_img_path + filename + ".jpg", os.getcwd() + '/images/value/' + str(filename) + ".jpg")
        if (filename.split('_')[0] == "anthrax"):
            shutil.copy(anthrax_img_path + filename + ".jpg", os.getcwd() + '/images/value/' + str(filename) + ".jpg")
        if (filename.split('_')[0] == "leafhopper"):
            shutil.copy(leafhopper_img_path + filename + ".jpg", os.getcwd() + '/images/value/' + str(filename) + ".jpg")

    txtlist = file_test.readlines()
    for filepath in txtlist:
        oneline = filepath.strip()
        filename = oneline.split('\\')[-1]
        shutil.copy(oneline + ".txt", os.getcwd() + '/labels/test/' + str(filename) + ".txt")

        if (filename.split('_')[0] == "normal"):
            shutil.copy(normal_img_path + filename + ".jpg", os.getcwd() + '/images/test/' + str(filename) + ".jpg")
        if (filename.split('_')[0] == "anthrax"):
            shutil.copy(anthrax_img_path + filename + ".jpg", os.getcwd() + '/images/test/' + str(filename) + ".jpg")
        if (filename.split('_')[0] == "leafhopper"):
            shutil.copy(leafhopper_img_path + filename + ".jpg", os.getcwd() + '/images/test/' + str(filename) + ".jpg")


def close_file():

    file_train.close()
    file_train_value.close()
    file_test.close()
    file_val.close()

if __name__ == "__main__":
    opt = get_opt()
    anthrax_text_path = opt.anthrax_text_path
    leafhopper_txt_path = opt.leafhopper_txt_path
    normal_txt_path = opt.normal_txt_path


    # create_txt_file(os.getcwd())
    # split_data(normal_txt_path, "normal")
    # split_data(anthrax_text_path, "anthrax")
    # split_data(leafhopper_txt_path, "leafhopper")
    # close_file()

    copy_flie(os.getcwd(), opt)
    close_file()
